import os
from pathlib import Path
import io

import streamlit as st
import torch
import torch.nn as nn
from PIL import Image
from torchvision import transforms as T
from torchvision.models import alexnet, AlexNet_Weights
import pandas as pd

# 路径与设备
BASE_DIR = Path(__file__).resolve().parent
# 数据目录解析：优先使用项目根目录的 natural_images，其次回退同级目录
DATA_DIR = next((
    p for p in [
        BASE_DIR.parent / "natural_images",  # 项目根目录外部共享数据集
        BASE_DIR / "natural_images",         # 回退：同级目录
    ] if p.exists()
), (BASE_DIR.parent / "natural_images"))
WEIGHT_PATH = BASE_DIR / "best_alexnet.pth"

def get_device(user_choice: str):
    if user_choice == "cuda" and torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")

# 与训练一致的预处理
def get_transform():
    return T.Compose([
        T.Resize((256, 256), antialias=True),
        T.CenterCrop((224, 224)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225])
    ])

# 类别名称（按照 ImageFolder 的字母序）
@st.cache_data(show_spinner=False)
def load_class_names():
    # 使用 ImageFolder 的类别读取方式，保证排序一致
    try:
        from torchvision.datasets import ImageFolder
        tmp = ImageFolder(str(DATA_DIR), transform=T.ToTensor())
        return tmp.classes
    except Exception:
        # 回退为已知类别顺序（natural_images 标注常见的8类）
        return ["airplane", "car", "cat", "dog", "flower", "fruit", "motorbike", "person"]

# 构建模型并加载权重
@st.cache_resource(show_spinner=False)
def load_model(num_classes: int, device: torch.device):
    model = alexnet(weights=AlexNet_Weights.IMAGENET1K_V1)
    # 冻结底层参数
    for p in model.parameters():
        p.requires_grad = False
    in_features = model.classifier[6].in_features
    model.classifier[6] = nn.Linear(in_features, num_classes)
    model.to(device)
    if WEIGHT_PATH.exists():
        state = torch.load(WEIGHT_PATH, map_location=device)
        model.load_state_dict(state, strict=False)
        model.eval()
        return model, True
    else:
        model.eval()
        return model, False

def infer_image(model, device, img: Image.Image, class_names):
    tfm = get_transform()
    x = tfm(img.convert("RGB")).unsqueeze(0).to(device)
    with torch.no_grad():
        logits = model(x)
        probs = logits.softmax(dim=-1).squeeze(0).cpu()
    top_prob, top_idx = probs.max(dim=0)
    pred_label = class_names[int(top_idx)]
    return pred_label, float(top_prob), probs.tolist()

def pick_sample_image():
    # 允许用户从数据集中选择一张样例图片
    if not DATA_DIR.exists():
        return None
    classes = sorted([p.name for p in DATA_DIR.iterdir() if p.is_dir()])
    if not classes:
        return None
    cls = st.selectbox("选择一个类别以从数据集中抽取样例", classes, index=0)
    img_dir = DATA_DIR / cls
    imgs = [p for p in img_dir.iterdir() if p.suffix.lower() in [".jpg", ".jpeg", ".png", ".bmp", ".webp"]]
    if not imgs:
        st.info("该类别下未找到图片")
        return None
    img_path = st.selectbox("选择样例图片", imgs, index=0)
    return img_path

def main():
    st.set_page_config(page_title="图像8分类（AlexNet+Streamlit）", page_icon="🖼️", layout="centered")
    st.title("图像8分类 Demo")
    st.caption("类别：airplane, car, cat, dog, flower, fruit, motorbike, person")

    # 侧边栏设置
    st.sidebar.header("设置")
    device_choice = st.sidebar.radio("推理设备", ["cpu", "cuda" if torch.cuda.is_available() else "cpu"], index=0)
    device = get_device(device_choice)
    class_names = load_class_names()
    model, has_weight = load_model(num_classes=len(class_names), device=device)
    if not has_weight:
        st.warning(f"未检测到已训练权重：{WEIGHT_PATH.name}。可先运行训练脚本训练模型。")

    tab1, tab2 = st.tabs(["上传图片", "选择样例"])
    with tab1:
        uploaded = st.file_uploader("上传一张图片进行分类", type=["jpg", "jpeg", "png", "bmp", "webp"])
        if uploaded:
            image = Image.open(io.BytesIO(uploaded.read()))
            st.image(image, caption="输入图片", use_container_width=True)
            if st.button("开始推理", type="primary"):
                pred, conf, probs = infer_image(model, device, image, class_names)
                st.success(f"预测：{pred}（置信度 {conf:.2%}）")
                df = pd.DataFrame({"class": class_names, "probability": probs})
                st.bar_chart(df, x="class", y="probability")

    with tab2:
        sample_path = pick_sample_image()
        if sample_path:
            image = Image.open(sample_path)
            st.image(image, caption=str(sample_path.name), use_container_width=True)
            if st.button("对样例图片推理"):
                pred, conf, probs = infer_image(model, device, image, class_names)
                st.success(f"预测：{pred}（置信度 {conf:.2%}）")
                df = pd.DataFrame({"class": class_names, "probability": probs})
                st.bar_chart(df, x="class", y="probability")

    st.divider()
    st.markdown(f"模型权重路径：`{WEIGHT_PATH}`")
    st.markdown(f"数据集路径：`{DATA_DIR}`")

if __name__ == "__main__":
    main()